// Copyright (c) 2025 刘颖
//  所属单位：中国科学院信息工程研究所网络空间安全防御重点实验室
//             Key Laboratory of Cyberspace Security Defense,Institute of Information Engineering, CAS
//             中国科学院大学网络空间安全学院
//             School of Cyber Security, University of Chinese Academy of Sciences     

// 本文件定义C语言版本512长NTT和基于NTT乘法实现代码，用于PolarLAC128
// This file implements the C-language version of 512-length NTT and NTT-based multiplication for PolarLAC128.

#include<stdint.h>
#include "ntt.h"
#include "polarlac_param.h"
#include "ntt_param.h"


#if defined LAC_LIGHT || defined LAC128
// This NTT function is a naive implementation of the Number Theoretic Transform (NTT) algorithm, 
// and we do not use it. It is provided for reference of the theorem of NTT algorithm.
// In-place NTT Algorithm for LAC128 with 512 length and 18433 as modulus
// Input: Coefficient vctor
// Output: NTT vector
int NTT(int32_t *a) 
{
    
    int t,m,i,j,s,e;
    int32_t U=0,V=0,S; // U and V store intermediate results, S stors the twiddle factor
    t=DIM_N; 
    // t is the step size, starting from DIM_N, 
    // t indics the distance between the elements to be processed in the same butterfly operation

    for(m=1;m<DIM_N;m<<=1) // Determines the current stage of the butterfly iteration
    {
        t=(t>>1);
        for(i=0;i<m;i++)  // Selects the appropriate root of unity for each butterfly operation
        {
            s=(i*t)<<1; // The START of the index of coefficients using the same twiddle factor
            e=s+t;  // The END of the index of coefficients using the same twiddle factor
            S=f[m+i];
            for(j=s;j<e;j++)  // Identifies the corresponding array elements to be processed
            {
                // CT Butterfly operation
                U=a[j];
                V=a[j+t]*S%NTTQ;
              
                a[j]=(U+V)%NTTQ;  // a[j] is in the range of (0, NTTQ)
                a[j+t]=(U+NTTQ-V)%NTTQ; // a[j+t] is in the range of (0, NTTQ)
            }
        }
    }
    
    return 0;
}

//The computation process of the INTT is analogous to that of the NTT, 
//and the two exhibit a symmetric relationship.
int INTT(int32_t *a) 
{
    
    int t,m,i,j,s,e,h;
    int32_t U=0,V=0,S; //same as ntt 
    t=1; //step
    //In INTT the step grows exponentially as a power of two.
   
    for(m=DIM_N;m>1;m>>=1) 
    {
        s=0; //start index
        h=m>>1;
        for(i=0;i<h;i++) 
        {
            e=s+t; // end index
            S=fn[h+i]; // INTT uses the inverse powers of roots as twiddle factors
            for(j=s;j<e;j++)  
            {
                // GS Butterfly operation
                U=a[j];
                V=a[j+t];

                a[j]=(U+V)%NTTQ; // Updates the first element of the butterfly

                a[j+t]=(U+NTTQ-V)*S%NTTQ;   // Updates the second element of the butterfly
            
            }
            s=s+2*t;// Moves to the next pair of coefficients
        }
        t*=2;// Doubles the step size at each stage
    }
    
    // The multiplication with 1/N mod NTTQ is needed in INTT
    for(i=0;i<DIM_N;i++)
    {
        a[i]=(a[i]*INVERSE_N)%NTTQ;
    }

    return 0;
}



int poly_mul_ntt(const uint16_t  *a, const uint16_t *s, uint16_t *b)
{
    int i;
    int32_t a_buf[DIM_N], s_buf[DIM_N];
    int32_t b_buf[DIM_N];

    for(i=0;i<DIM_N;i++)
    {
        a_buf[i]=a[i];
        s_buf[i]=s[i];
    }
   
     //NTT form
    NTT(a_buf);  
    NTT(s_buf);  
    
     //point mul
     for(i=0;i<DIM_N;i++)
     {
         b_buf[i]=(a_buf[i]%NTTQ)*s_buf[i]%NTTQ;  
     }

     //INTT to recover the coefficient form
    INTT(b_buf);

     //send back the results
     for(i=0;i<DIM_N;i++)
     {
         b[i]=b_buf[i];
     }
     
     return 0;
}



/**
 * @brief Computes the Montgomery map of the input value.
 *
 * This function calculates r = a * β^(-1) mod NTTQ, where 0 < r < NTTQ
 * and maps the result to its standard representation.
 *
 * @param a The input integer value to be mapped.
 * @return The mapped integer value r in the range (0, NTTQ).
 */
static int32_t MontgomerymapFull(int32_t a)  
{
    int32_t t;
    int32_t m,r;
    m = a * INVERSE_Q;
    m=m&((1<<16)-1);  // m mod β
    t = m * NTTQ;
    r = a - t;
    r >>= 16; // r in thr range of (-NTTQ, NTTQ)
    r = r + ((r >> 15) & NTTQ); // Ensure r is in the range of (0, NTTQ)
    return r;
}


/**
 * @brief Performs the Number Theoretic Transform (NTT) with lazy reduction.
 * 
 * This function applies the NTT algorithm to the input array `a` using lazy 
 * modular reduction. Lazy reduction allows intermediate values to exceed 
 * the modulus temporarily, reducing the number of modular operations and 
 * improving performance.
 * 
 * @param a Pointer to the input array of integers (int16_t) to be transformed.
 *          The array is modified in-place to contain the transformed values.
 * 
 * @return Returns an integer status code. Typically, 0 indicates success, 
 *         while other values may indicate errors or specific conditions.
 */
int NTT_Lazy(int16_t *a)
{
    int t,m,i,j,s,e;
    int16_t U=0,V=0; // U and V store intermediate results in butterfly operations
    uint16_t S;  // S stores the twiddle factor for the current butterfly operation
    t=DIM_N;
    // t is the step size, starting from DIM_N, 
    // t indics the distance between the elements to be processed in the same butterfly operation


    // First level in the computation in NTT
    m=1;
    t=(t>>1);// Halve the step size
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;     // Start index for coefficients using the same twiddle factor
            e=s+t;         // End index for coefficients using the same twiddle factor
            S=M[m+i];     // Twiddle factor for the current stage
            //Since the output of the Montgomery algorithm is rβ^(-1) mod NTTQ, 
            // a scaling factor β is introduced into the root-of-unity array M to eliminate the Montgomery factor. 
            // As a result, the roots of unity used here differ from those f[i] in the standard NTT algorithm.
            // That is M[i]=f[i]*β mod NTTQ

            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S); // Use Montgomery to reduce multiplications

                // Because the all inputs at this stage are within (0, Q), 
                // modular reduction is not required for addition and subtraction.
                a[j]=U+V-NTTQ; // in the range of [-NTTQ, NTTQ]
                a[j+t]=U-V; // in the range of [-NTTQ, NTTQ]
              
            }
        }

    // Subsequent stages of the butterfly operation (m = 2, 4, ..., DIM_N/2)
    // 2~log_2N−2 level : perform the modular reduction to U 
    for(m=2;m<DIM_N_2;m<<=1) 
    {
        t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=M[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S);
                U = U + ((U>>15)&NTTQ); // lightweight modular reduction to the U to map it to （-NTTQ, NTTQ）

                a[j]=(U+V-NTTQ);
                a[j+t]=(U-V);
              
            }
        }
    }

    //final level : perform the modular reductions to U, a[j] and a[j+t]  
    m=DIM_N_2;
    t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=M[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S);
                U = U + ((U>>15)&NTTQ);
                
                //Map each coefficient to its standard representation.
                a[j]=U+V-NTTQ;
                a[j] = a[j] + ((a[j]>>15)&NTTQ);

                a[j+t]=U-V;
                a[j+t] = a[j+t] + ((a[j+t]>>15)&NTTQ);
            }
        }


    return 0;
}



/**
 * @brief Performs the Inverse Number Theoretic Transform (INTT) with lazy reduction.
 * 
 * This function applies the INTT algorithm to the input array `a` using lazy 
 * modular reduction. Lazy reduction allows intermediate values to exceed 
 * the modulus temporarily, reducing the number of modular operations and 
 * improving performance.
 * 
 * @param a Pointer to the input array of integers (int16_t) to be transformed.
 *          The array is modified in-place to contain the transformed values.
 * 
 * @return Returns an integer status code. Typically, 0 indicates success, 
 *         while other values may indicate errors or specific conditions.
 */

/**
 *Notice： 
 * 
 * GS butterfly structure used in the inverse transform exhibits weaker symmetry properties. 
 * Due to the use of a 15-bit modulus, the intermediate results of additions and subtractions 
 * are more likely to exceed the modulus range. 
 * So modular reduction must be performed explicitly at each layer of the inverse transform.
 * 
 */
int INTT_Lazy(int16_t *a)
{
    int t, m, i, j, s, e, h;
    int16_t U = 0, V = 0; // U and V store intermediate results in butterfly operations
    uint16_t S;           // S stores the twiddle factor for the current butterfly operation
    t = 1;                // Step size, starting from 1 and doubling at each stage

    // Perform the butterfly operations for each stage of the INTT
    for (m = DIM_N; m > 1; m /= 2) 
    {
        s = 0;           // Start index for coefficients using the same twiddle factor
        h = m / 2;       // Half the size of the current stage
        for (i = 0; i < h; i++) 
        {
            e = s + t;   // End index for coefficients using the same twiddle factor
            S = Mn[h + i]; // Twiddle factor for the current stage
            for (j = s; j < e; j++) 
            {
                U = a[j];
                V = a[j + t];

                // Update the first element of the butterfly
                a[j] = (U + V - NTTQ);
                a[j] = a[j] + ((a[j] >> 15) & NTTQ); // Ensure the result is in the range (0, NTTQ)

                // Update the second element of the butterfly using Montgomery reduction
                a[j + t] = MontgomerymapFull((int32_t)(U - V) * (int32_t)S);
            }
            s = s + 2 * t; // Move to the next pair of coefficients
        }
        t *= 2; // Double the step size at each stage
    }

    // Multiply each coefficient by the scaling factor N^(-1)*β mod NTTQ
    // The β is used to emliminate the Montgomery factor introduced in the point-mul stage in the fucntion poly_mul_ntt_Lazy
    for (i = 0; i < DIM_N; i++) 
    {
        a[i] = (a[i] * INVERSE_N_BETA) % NTTQ;
    }

    return 0;
}


/**
 * Performs polynomial multiplication using the Number Theoretic Transform (NTT) in a lazy manner.
 *
 * This function multiplies two polynomials represented as arrays of 16-bit unsigned integers.
 * The result is stored in the output array `b`. The multiplication is performed using NTT,
 * which is efficient for large polynomial operations.
 *
 * @param a Pointer to the first polynomial array (input).
 * @param s Pointer to the second polynomial array (input).
 * @param b Pointer to the output polynomial array where the result will be stored.
 * @return An integer indicating the success or failure of the operation.
 */
int poly_mul_ntt_Lazy(const uint16_t  *a, const uint16_t *s, uint16_t *b)
{
    int i;
    int16_t a_buf[DIM_N], s_buf[DIM_N],b_buf[DIM_N];

    for(i=0;i<DIM_N;i++)
    {
        a_buf[i]=a[i];
        s_buf[i]=s[i];
    }
   
     //NTT form
    NTT_Lazy(a_buf);
    NTT_Lazy(s_buf);
    
     //point mul
     for(i=0;i<DIM_N;i++)
     {
         b_buf[i]=MontgomerymapFull((int32_t)a_buf[i]*(int32_t)s_buf[i]);
         // There will introduced a Montgomery factor β^(-1) mod NTTQ
         // So we perform the multiplication with N^(-1)*β mod NTTQ in the final step of INTT
     }

     //INTT form
    INTT_Lazy(b_buf);

     //mod Q
     for(i=0;i<DIM_N;i++)
     {
         b[i]=(b_buf[i]+NTTQ)%NTTQ; // To ensure the results are all positive numbers.
     }

     return 0;
}
#endif